# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time        : 2020/12/16 10:33
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : iteration_methods.py
@Version     : Version 1.0.0
@Description : TODO 自己实现雅各比迭代法，高斯-赛德尔迭代法,SOR迭代法，谱半径函数，希尔伯特矩阵函数
@Created By  : PyCharm
"""

import numpy as np
from scipy.linalg import eigvals, inv

import \
    AI.NumericalAnalysis.chapter5_directive_solve_linear_system_of_equations.DirectiveSolveLinearSystemEquations as dsl


def vector_norm(vector: np.ndarray, p=None):
    """
    计算向量的p-范数
    :param vector: 实向量或者复向量
    :param p: 指定类型的范数，默认是oo范数
    :return: 指定向量范数和p值
    """
    if p is None:
        return abs(vector).max(), p
    elif p >= 1:
        return np.power(np.sum(np.power(abs(vector), p)), 1 / p), p
    else:
        raise Exception("error,p must be an integer , greater than  or equal to 1")


def jacobi(coefficient_matrix: np.ndarray,
           right_hand_side_vector: np.ndarray,
           initial_vector: np.ndarray, epsilon=1e-4):
    d = np.diag(np.diag(coefficient_matrix))
    l = -np.tril(coefficient_matrix, k=-1)
    u = -np.triu(coefficient_matrix, k=1)
    jacobi_matrix = inv(d) @ (l + u)
    if spectral_radius(jacobi_matrix) < 1:
        f = inv(d) @ right_hand_side_vector
        x_1 = initial_vector
        x_2 = jacobi_matrix @ x_1 + f
        iteration_count = 1
        while 1:
            norm, _ = vector_norm(x_2 - x_1)
            if norm < epsilon:
                return x_2, iteration_count
            x_1 = x_2
            x_2 = jacobi_matrix @ x_1 + f
            iteration_count += 1
    raise Exception("jacobi iteration is not convergent")


def gauss_seidel(coefficient_matrix: np.ndarray,
                 right_hand_side_vector: np.ndarray,
                 initial_vector: np.ndarray, epsilon=1e-4):
    d = np.diag(np.diag(coefficient_matrix))
    l = -np.tril(coefficient_matrix, k=-1)
    u = -np.triu(coefficient_matrix, k=1)
    gauss_seidel_matrix = inv(d - l) @ u
    if spectral_radius(gauss_seidel_matrix) < 1:
        f = inv(d - l) @ right_hand_side_vector
        x_1 = initial_vector
        x_2 = gauss_seidel_matrix @ x_1 + f
        iteration_count = 1
        while 1:
            norm, _ = vector_norm(x_2 - x_1)
            if norm < epsilon:
                return x_2, iteration_count
            x_1 = x_2
            x_2 = gauss_seidel_matrix @ x_1 + f
            iteration_count += 1
    raise Exception("jacobi iteration is not convergent")


def successive_over_relaxation(coefficient_matrix: np.ndarray,
                               right_hand_side_vector: np.ndarray,
                               initial_vector: np.ndarray,
                               true_solution: np.ndarray, omega=1.0, epsilon=5e-6):
    d = np.diag(np.diag(coefficient_matrix))
    l = -np.tril(coefficient_matrix, k=-1)
    u = -np.triu(coefficient_matrix, k=1)
    l_omega = inv(d - omega * l) @ ((1 - omega) * d + omega * u)
    if spectral_radius(l_omega) < 1:
        f = omega * inv(d - omega * l) @ right_hand_side_vector
        x = l_omega @ initial_vector + f
        iteration_count = 1
        while 1:
            norm, _ = vector_norm(true_solution - x)
            if norm < epsilon:
                return x, iteration_count
            x = l_omega @ x + f
            iteration_count += 1
    raise Exception("jacobi iteration is not convergent")


def hilbert_matrix(order: int):
    hilbert = np.zeros((order, order))
    for i in range(order):
        for j in range(order):
            hilbert[i, j] = 1 / (i + j + 1)
    return hilbert


def spectral_radius(square_matrix: np.ndarray):
    if square_matrix.shape[0] == square_matrix.shape[1]:
        return abs(eigvals(square_matrix)).max()
    raise Exception("\n{} is not a square matrix".format(square_matrix))


if __name__ == '__main__':
    # 第一组测试方程组，来源详见李庆扬数值分析第5版P180-181
    # A = np.array([[8, -3, 2], [4, 11, -1], [6, 3, 12]], dtype=np.float64)
    # b = np.array([20, 33, 36], dtype=np.float64).reshape((3, 1))
    # true_solution_vector = dsl.gaussian_elimination(A, b)
    # b0 = np.array([[0, 3 / 8, -2 / 8], [-4 / 11, 0, 1 / 11], [-6 / 12, -3 / 12, 0]])
    # print(spectral_radius(b0))

    # 第二组测试方程组，来源详见李庆扬数值分析第5版P209 e.g.1
    """
    A = np.array([[5, 2, 1], [-1, 4, 2], [2, -3, 10]], dtype=np.float64)
    b = np.array([-12, 20, 3], dtype=np.float64).reshape((3, 1))
    x0 = np.array([1, 1, 1], dtype=np.float64).reshape((3, 1))
    root1, count1 = jacobi(A, b, x0)
    print("jacobi公式迭代{}次达到精度要求，\n近似解x=\n{}".format(count1, root1))
    root2, count2 = gauss_seidel(A, b, x0)
    print("Gauss_seidel公式迭代{}次达到精度要求，\n近似解x=\n{}".format(count2, root2))
    """
    # 第三组测试方程组，来源详见李庆扬数值分析第5版P210 e.g.7
    """
    A = np.array([[4, -1, 0], [-1, 4, -1], [0, -1, 4]], dtype=np.float64)
    b = np.array([1, 4, -3], dtype=np.float64).reshape((3, 1))
    true_roots = np.array([1 / 2, 1, -1 / 2]).reshape((3, 1))
    x0 = np.array([0, 0, 0]).reshape((3, 1))
    # omega=1.03
    root1, count1 = successive_over_relaxation(A, b, x0, true_roots, 1.03)
    print("SOR方法,当取Omega=1.03,初值x0=\n{}\n迭代{}次可达到精度要求，此时近似解x=\n{}".format(x0, count1, root1))
    # omega=1
    root2, count2 = successive_over_relaxation(A, b, x0, true_roots)
    print("SOR方法,当取Omega=1,初值x0=\n{}\n迭代{}次可达到精度要求，此时近似解x=\n{}".format(x0, count2, root2))
    # omega=1.1
    root3, count3 = successive_over_relaxation(A, b, x0, true_roots, 1.1)
    print("SOR方法,当取Omega=1.1,初值x0=\n{}\n迭代{}次可达到精度要求，此时近似解x=\n{}".format(x0, count3, root3))
    """
    # 第四组测试方程组，来源详见李庆扬数值分析第5版P211 e.g.1
    A = hilbert_matrix(10)
    b = A @ np.ones((10, 1))
    x0 = np.zeros((10, 1))
    true_roots = dsl.gaussian_elimination(A, b)
    print("列主元高斯消去法精确解x=\n{}".format(true_roots))
    # root1, count1 = jacobi(A, b, x0, 1e-6)
    # print("jacobi公式迭代{}次达到精度要求，\n近似解x=\n{}".format(count1, root1))
    # omega=1
    root2, count2 = successive_over_relaxation(A, b, x0, true_roots)
    print("SOR方法,当取Omega=1,初值x0=\n{}\n迭代{}次可达到精度要求，此时近似解x=\n{}".format(x0, count2, root2))
    # omega=1.25
    root3, count3 = successive_over_relaxation(A, b, x0, true_roots, 1.25)
    print("SOR方法,当取Omega=1.25,初值x0=\n{}\n迭代{}次可达到精度要求，此时近似解x=\n{}".format(x0, count3, root3))
    # omega=1.5
    root4, count4 = successive_over_relaxation(A, b, x0, true_roots, 1.5)
    print("SOR方法,当取Omega=1.5,初值x0=\n{}\n迭代{}次可达到精度要求，此时近似解x=\n{}".format(x0, count4, root4))
